import argparse
import os, shutil, math, time
import subprocess
import ffmpeg
import tqdm

root_dir = os.path.dirname(__file__)
srmd_cli = os.path.join(root_dir, "srmd-ncnn-vulkan-20220728-ubuntu/srmd-ncnn-vulkan")
rife_cli = os.path.join(root_dir, "rife-ncnn-vulkan-20221029-ubuntu/rife-ncnn-vulkan")

type_4k = {
    0 : [4096, 3112], #Full Aperture 4K
    1: [3656, 2664], #Academy 4K
    2: [3840, 2160], #UHDTV
}

def probe_video_info(input_video):
    probe = ffmpeg.probe(input_video)
    video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
    if not video_stream:
        raise Exception("No video stream found")
    duration = float(probe['format']['duration'])
    width = int(video_stream['width'])
    height = int(video_stream['height'])    
    num_frames = int(video_stream['nb_frames'])
    fps = float(video_stream['avg_frame_rate'].split('/')[0]) / float(video_stream['avg_frame_rate'].split('/')[1])
    video_codec = video_stream['codec_name']
    video_info = {
        'duration': duration,
        'width': width,
        'height': height,
        'num_frames': num_frames,
        'fps': fps,
        'video_codec': video_codec,
    }
    audio_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)
    if audio_stream:
        audio_codec = audio_stream['codec_name']
        video_info['audio_codec'] = audio_codec

    print(video_info)
    return video_info
def extract_frames(input_video, output_dir):
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)
    
    print(f"Extracting frames to {output_dir}...")
    ffmpeg_proc = subprocess.Popen(["ffmpeg",
        "-i", input_video,
        "-q:v", "1",
        os.path.join(output_dir, "%04d.png")],
        # stdin=subprocess.PIPE,
        # stdout=subprocess.PIPE,
        # stderr=subprocess.PIPE
        )
    
    # while True:
    #     line = ffmpeg_proc.stderr.readline()
    #     if line:
    #         line = line.decode("utf-8")
    #         print(line)
    #     elif ffmpeg_proc.poll() is not None:
    #         break       
    ffmpeg_proc.wait()
    print("Frame Extraction Done.")

def extract_audio(input_video, output_audio):
    print("Extracting audio...")
    if os.path.exists(output_audio):
        os.remove(output_audio)

    ffmpeg_proc = subprocess.Popen(["ffmpeg",
        "-i", input_video,
        "-map", "0:a",
        "-acodec", "copy",
        output_audio],
        # stdin=subprocess.PIPE,
        # stdout=subprocess.PIPE,
        # stderr=subprocess.PIPE
        )
    
    # while True:
    #     line = ffmpeg_proc.stderr.readline()        
    #     if line:
    #         line = line.decode("utf-8")
    #         print(line)
    #     elif ffmpeg_proc.poll() is not None:
    #         break

    ffmpeg_proc.wait()
    print("Audio Extraction Done.")    

def interpolate_frames(frames_dir, output_dir, target_num_frames, model="rife-v4.6"):
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)

    print(f"Inserting frames to {target_num_frames}...")
    rife_proc = subprocess.Popen([rife_cli,
                    "-i", frames_dir,
                    "-o", output_dir,
                    "-m", model,
                    "-n", str(target_num_frames),
                    "-j", "4:8:8"
                    ],
                    # stdin=subprocess.PIPE,
                    # stdout=subprocess.PIPE,
                    # stderr=subprocess.PIPE
                    )    
    # while True:
    #     line = rife_proc.stderr.readline()
    #     if line:
    #         line = line.decode("utf-8")
    #         print(line)
    #     elif rife_proc.poll() is not None:
    #         break

    progress = tqdm.trange(int(target_num_frames), desc="interpolating frames")
    while rife_proc.poll() is None:
        time.sleep(2)
        num_files = len(os.listdir(output_dir))
        if num_files > progress.n:
            progress.update(num_files - progress.n)

    if rife_proc.returncode != 0:
        raise Exception("RIFE failed")
    
    num_files = len(os.listdir(output_dir))
    progress.update(target_num_frames - progress.n)
    progress.close()
    print("Interpolation Done.")

def calculate_scales(target_scale, scale_per_time):
    if target_scale <= scale_per_time:
        return [target_scale]
    else:
        scales = [scale_per_time, target_scale // scale_per_time, target_scale % scale_per_time]
        while scales[-1] > 0:
            if scales[-2] >= scale_per_time:                
                scales[-1] = scales[-2] % scale_per_time
                scales[-2] = scale_per_time
            else:
                scales[-2] += 1                
                scales.pop()
                break

        return scales
def upscale_frames(frames_dir, output_dir, scale, noise_level=-1):
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
        
    os.makedirs(output_dir)
    
    print(f"Upscaling frames to {scale} times with noise level {noise_level}...")    
    srmd_proc = subprocess.Popen([srmd_cli,
                    "-i", frames_dir,
                    "-o", output_dir,
                    "-s", str(scale),
                    "-n", str(noise_level-1),
                    "-t", "1024",
                    "-j", "2:4:4"],
                    # stdin=subprocess.PIPE,
                    # stdout=subprocess.PIPE,
                    # stderr=subprocess.PIPE
                    )
    
    target_num_frames = len(os.listdir(frames_dir))
    progress = tqdm.trange(int(target_num_frames), desc="upscale frames")
    while srmd_proc.poll() is None:
        time.sleep(2)
        num_files = len(os.listdir(output_dir))
        if num_files > progress.n:
            progress.update(num_files - progress.n)
    
    if srmd_proc.returncode != 0:
        raise Exception("SRMD failed")
    
    num_files = len(os.listdir(output_dir))
    progress.update(target_num_frames - progress.n)
    progress.close()

def merge_frames(frames_dir, audio, target_video, target_width, target_height, target_fps):
    if os.path.exists(target_video):
        os.remove(target_video)

    frame_files = os.listdir(frames_dir)
    num_digits = len(os.path.splitext(frame_files[0])[0])
    print("Merging frames...")
    ffmpeg_proc = subprocess.Popen(["ffmpeg",
        "-f", "image2",
        "-i", os.path.join(frames_dir, f"%0{num_digits}d.png"),
        "-i", audio,
        "-r", str(target_fps),
        "-s", f"{target_width}x{target_height}",
        "-vcodec", "libx264",
        "-acodec", "copy",
        "-pix_fmt", "yuv420p",
        "-crf", "15",
       target_video],
    #    stdin=subprocess.PIPE,
    #    stdout=subprocess.PIPE,
    #    stderr=subprocess.PIPE
    )
    
    # while True:
    #     line = ffmpeg_proc.stderr.readline()
    #     if line:
    #         line = line.decode("utf-8")
    #         print(line)
    #     elif ffmpeg_proc.poll() is not None:
    #         break
    ffmpeg_proc.wait()
    print("Merge Done.")

def main(input, output, res, fps, noise, step, recover):
    width, height = type_4k[res]

    video_info = probe_video_info(input)
    work_dir = os.path.dirname(output)
    log_file = os.path.join(work_dir, "log.txt")
    steps = ["Audio Extraction Done.", "Frame Extraction Done.", "Interpolation Done.", "Upscale Done.", "Merge Done."]
    if recover and os.path.exists(log_file):
        with open(log_file, "r") as f:
            done_steps = [line.strip() for line in f.readlines() if line.strip() in steps]
    else:
        done_steps = []

    try:
        audio_path = os.path.join(work_dir, "audio" + os.path.splitext(input)[1])
        if steps[0] not in done_steps:
            extract_audio(input, audio_path)
            done_steps.append(steps[0])
        
        raw_frames_dir = os.path.join(work_dir, "raw_frames")
        if steps[1] not in done_steps:
            extract_frames(input, raw_frames_dir)
            done_steps.append(steps[1])
        
        interpolate_frames_dir = os.path.join(work_dir, "insert_frames")    
        if steps[2] not in done_steps:
            num_target_frames = int(math.ceil(len(os.listdir(raw_frames_dir)) * fps / video_info['fps']))
            interpolate_frames(raw_frames_dir, interpolate_frames_dir, num_target_frames)
            done_steps.append(steps[2])
        
        interpolated_frame_files = [file for root, dirs, files in os.walk(interpolate_frames_dir) for file in files if file.endswith(".png")]
        interpolated_frame_files.sort()
        num_steps = int(math.ceil(len(interpolated_frame_files) / fps / step))
        if steps[3] not in done_steps:
            upscale = int(math.ceil(min(width / video_info['width'], height / video_info['height'])))
            scales = calculate_scales(upscale, 2)
            scales3 = calculate_scales(upscale, 3)
            scales4 = calculate_scales(upscale, 4)
            if len(scales3) < len(scales):
                scales = scales3

            if len(scales4) < len(scales):
                scales = scales4

            print(f"Upscaling frames to {upscale}x by iteration series {scales}...")
            
            num_frames_per_step = int(step * fps)
            if len(interpolated_frame_files) and os.path.exists(os.path.join(interpolate_frames_dir, interpolated_frame_files[0])) > 0:
                for i in range(num_steps):
                    os.makedirs(os.path.join(interpolate_frames_dir, f'{i}'))
                    for j in range(i * num_frames_per_step, min((i + 1) * num_frames_per_step, len(interpolated_frame_files))):
                        shutil.move(os.path.join(interpolate_frames_dir, interpolated_frame_files[j]), os.path.join(interpolate_frames_dir, f'{i}'))
            
            for i in range(len(scales)):     
                print(f"Upscaling step {i} with scale {scales[i]}...")
                if i == 0:
                    input_frames_dir = interpolate_frames_dir
                    upscale_frames_dir = os.path.join(work_dir, f"upscale_frames_{i}")
                else:            
                    if i > 1:
                        shutil.rmtree(input_frames_dir)

                    input_frames_dir = upscale_frames_dir
                    upscale_frames_dir = os.path.join(work_dir, f"upscale_frames_{i}")

                if recover and os.path.exists(upscale_frames_dir):
                    step_dirs = [folder for folder in os.listdir(upscale_frames_dir) if os.path.isdir(os.path.join(upscale_frames_dir, folder))]
                    if len(step_dirs) == num_steps:
                        num_frames_in_last_step = len(os.listdir(os.path.join(upscale_frames_dir, step_dirs[-1])))
                        num_frames_in_last_step_to_handle = len(os.listdir(os.path.join(input_frames_dir, step_dirs[-1])))
                        if num_frames_in_last_step == num_frames_in_last_step_to_handle:
                            print("Skip upscaling step since it is already done.")
                            continue                        
                    
                for s in tqdm.trange(num_steps, desc="Upscaling splitted steps"):                    
                    src_dir = os.path.join(input_frames_dir, f'{s}')
                    target_dir = os.path.join(upscale_frames_dir, f'{s}')
                    num_frames_to_upscale = len(os.listdir(src_dir))
                    if not recover or not os.path.exists(target_dir) or num_frames_to_upscale > len(os.listdir(target_dir)):
                        print(f"Upscaling step-{s} from {s * step} secs to {s * step + num_frames_to_upscale // fps} secs...")
                        upscale_frames(src_dir,
                                target_dir,
                                scales[i], noise)
                    else:
                        print(f"Skip upscaling step {s} since it is already done.")

            if len(scales) > 1:
                shutil.rmtree(input_frames_dir)

            done_steps.append(steps[3])
            print("Upscale Done.")
        else:            
            upscale_frames_dir = [dir for dir in os.listdir(work_dir) if dir.startswith("upscale_frames")]
            if len(upscale_frames_dir) == 1:
                upscale_frames_dir = os.path.join(work_dir, upscale_frames_dir[0])
            else:
                upscale_frames_dir = None
        
        if steps[4] not in done_steps:
            merge_frames_dir = os.path.join(work_dir, "merge_frames")
            if not os.path.exists(merge_frames_dir):
                os.makedirs(merge_frames_dir)
            if os.path.exists(upscale_frames_dir):
                step_dirs = [os.path.join(upscale_frames_dir, f'{i}') for i in range(num_steps)]                
                for step_dir in step_dirs:
                    for frame_file in os.listdir(step_dir):
                        if frame_file.endswith(".png"):
                            target_frame_file = os.path.join(merge_frames_dir, frame_file)
                            if os.path.exists(target_frame_file):
                                os.remove(target_frame_file)
                            shutil.move(os.path.join(step_dir, frame_file), target_frame_file)

                    os.rmdir(step_dir)
                os.rmdir(upscale_frames_dir)
            
            merge_frames(merge_frames_dir, audio_path, output, width, height, fps)
            done_steps.append(steps[4])

        if os.path.exists(output):
            elapsed = int(time.time() - start)
            seconds = elapsed % 60
            minutes = elapsed // 60 % 60
            hours = elapsed // 3600
            print(f"Upscale video succeed to {output}, {hours}:{minutes}:{seconds} time used.")
            probe_video_info(output)
        else:
            print("Upscale video failed, check the log for errors.")
    finally:
        with open(log_file, "w") as f:            
            f.write(str(args))
            f.write("\n")
            f.write("Progress:\n")
            for step in done_steps:
                f.write(step + "\n")


if __name__ == "__main__":
    start = time.time()
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", type=str, help="input video")
    parser.add_argument("-o", "--output", type=str, help="output video")
    parser.add_argument("-r", "--res", type=int, help="resolution type of the video, 0 for Full Aperture 4k, 1 for Academy 4k, 2 for UHDTV 4k", default=0)
    parser.add_argument("-f", "--fps", type=int, help="frame rate", default=60)
    parser.add_argument("-n", "--noise", type=int, help="noise level, from 0 to 11", default=3)
    parser.add_argument("-s", "--step", type=int, help="processing step duration in seconds", default=30)
    parser.add_argument("-R", "--recover", type=bool, help="recover processing from previous saved log", default=True)

    args = parser.parse_args()
    main(args.input, args.output, args.res, args.fps, args.noise, args.step, args.recover)
    
    
    